// kdtree_template.h
// 04/21/2011 : created
#pragma once
#include <vector>
#include <QtCore/QVector>
#include <QtCore/QtAlgorithms>
#include "arma_vector.h"
#include "kdtree.h"
#include "tree.h"

template <class DerivedTree>
class KDTreeTemplate : public Tree
{
protected:
  typedef Tree::dataset_type    dataset_type;
  typedef std::vector<int>      vint_type;
protected:
  BoundingBox m_box;
  vint_type& m_dimList;
public:
  KDTreeTemplate(dataset_type& data) : Tree(data), m_box(n_dim()), m_dimList(*(new vint_type(n_dim()))) { 
    random_permutation(m_dimList); 
  }
  KDTreeTemplate(KDTreeTemplate* parent, int s, int e) : Tree(parent, s, e), m_box(n_dim()), m_dimList(parent->m_dimList)
  {
  }
  ~KDTreeTemplate()
  {
    if (!m_parent) {
      delete &m_dimList;
    }
  }

  void buildTree()
  {
    int n_left = -1;

    // find dimension with max variance
    double maxVariance = 0, maxDim = -1;

    if (n() > (int) m_data.size()/100 && n() > 1) {  // only split big nodes
      for (int dim = 0; dim < n_dim(); dim++) {
        double min = POSITIVE_INFINITY, max = -min;
        for_all_node_data(i) {
          double v = get(i)[dim];
          if (min > v) min = v;
          if (max < v) max = v;
        }
        double variance = max-min;
        if (maxVariance < variance) {
          maxVariance = variance;
          maxDim = dim;
        }
      }
    }
      
    if (maxDim != -1) {
      n_left = split(maxDim);
    }

    //if (n() > (int) m_data.size()/100 && n() > 1) {  // only split big nodes
    //  int dimIdx = rand() % n_dim();
    //  //Log::m_log << "split node n=" << n() << endl;
    //  for (int dim = 0; dim < n_dim(); dim++, dimIdx = (dimIdx+1)%n_dim()) {
    //    n_left = split(m_dimList[dimIdx]);
    //    //Log::m_log << "try split at dim=" << m_dimList[dimIdx] << " n_left=" << n_left << endl;
    //    if (n_left != -1) {
    //      //Log::m_log << "split at dim=" << m_dimList[dimIdx] << endl;
    //      break;
    //    }
    //  }
    //}

    if (n_left != -1) {
      // split
      KDTreeTemplate *left = new DerivedTree(static_cast<DerivedTree*>(this), m_s, m_s+n_left);
      KDTreeTemplate *right = new DerivedTree(static_cast<DerivedTree*>(this), m_s+n_left, m_s+n());
      m_child.push_back(left);
      m_child.push_back(right);
      left->buildTree();
      right->buildTree();
      // build node stats
      m_box.accumulate(left->m_box);
      m_box.accumulate(right->m_box);
    }
    else { // build leaf stats
      for_all_node_data(i) {
        m_box.accumulate(get(i));
      }
    }
  }
  const BoundingBox& box() const { return m_box; }
  double findKNN(const vec& query, std::vector<int>& knnIdx, std::vector<double>& minBound) {
    return static_cast<DerivedTree*>(this)->implementKNN(query, knnIdx, minBound);
  }
  DerivedTree* child(int i) const
  {
    assert(i >= 0 && i < (int) m_child.size());
    return static_cast<DerivedTree*>(this->m_child[i]);
  }
protected:
  int split(int dim)
  {
    // choose split value in dimension dim
    QVector<double> data(n());
    for_all_node_data(i) data[i] = get(i)[dim];
    qSort(data.begin(), data.end());

    double splitVal = data[n()/2];
    if (n()/2+1 < n()) splitVal = (data[n()/2] + data[n()/2+1]) / 2;

    // try splitting the node by moving indices with value < splitVal to the left
    int n_left = 0;
    for_all_node_data(i) if (get(i)[dim] < splitVal) {
      int tmp = m_idx[m_s+n_left]; m_idx[m_s+n_left] = m_idx[m_s+i]; m_idx[m_s+i] = tmp;
      n_left++;
    }
    if (n_left > 0 && n_left < n())
      return n_left;  // success, return number of data points in left node
    else
      return -1; // failure
  }

  static void random_permutation(std::vector<int>& p)
  {
    int n = p.size();
    for (int i = 0; i < n; i++) p[i] = i;
    for (int i = 0; i < n-1; i++) {
      int swap = rand() % (n-i) + i;
      int tmp = p[i]; p[i] = p[swap]; p[swap] = tmp;
    }
  }
};
